#include <monocular/Monocular.hpp>

namespace wmj
{
    Monocular::Monocular(std::string name) : Node(name)
    {
        setParam(MONO_CFG);

        // 创建回调组，多线程接收msg
        rclcpp::CallbackGroup::SharedPtr callback_group_armors = this->create_callback_group(rclcpp::CallbackGroupType::MutuallyExclusive);
        auto sub_armors_opt = rclcpp::SubscriptionOptions();
        sub_armors_opt.callback_group = callback_group_armors;
        m_armors_sub = this->create_subscription<base_interfaces::msg::Armors>("/Armors", rclcpp::QoS(4), std::bind(&Monocular::armor_cb, this, std::placeholders::_1), sub_armors_opt);

        m_armors_pub = this->create_publisher<base_interfaces::msg::Armors>("/Armors_base",1);
    }

    void Monocular::setParam(const std::string &file_path)
    {
        cv::FileStorage fs(file_path, cv::FileStorage::READ);
        fs["debug"] >> m_debug;
        fs["ArmorParams"]["small_armor_height"] >> m_small_height;
        fs["ArmorParams"]["small_armor_width"] >> m_small_width;
        fs["ArmorParams"]["large_armor_height"] >> m_large_height;
        fs["ArmorParams"]["large_armor_width"] >> m_large_width;

        // 获取装甲板世界坐标系的坐标点
        getObjPoints(m_small_width, m_small_height, wmj::ARMORTYPE::ARMOR_SMALL);
        getObjPoints(m_large_width, m_large_height, wmj::ARMORTYPE::ARMOR_LARGE);

        // 读取相机参数
        for (auto node: fs["Cameras"])
        {
            std::string frame_id;
            cv::Mat CameraMat, DistMat;
            node["frame_id"] >> frame_id;
            node["camera_matrix"] >> CameraMat;
            node["distortion_coefficients"] >> DistMat;
            this->m_camera_frame_id.push_back(frame_id);
            this->m_CameraMat.push_back(CameraMat);
            this->m_DistMat.push_back(DistMat);
        }
        fs.release();
    }

    void Monocular::armor_cb(const base_interfaces::msg::Armors::SharedPtr msg)
    {
        // 创建装甲板消息
        base_interfaces::msg::Armors armors_msg;
        armors_msg.num = 0;
        if (msg->header.stamp.sec != 0 && msg->header.stamp.nanosec != 0)
        {
            armors_msg.header.stamp = msg->header.stamp;
        }
        else 
        {
            armors_msg.header.stamp = this->now();
        }

        // 设置坐标系
        if (msg->header.frame_id != "") //假如frame_id不为空
        {
            armors_msg.header.frame_id = msg->header.frame_id;
        }
        else if (msg->num > 0)
        {
            armors_msg.header.frame_id = msg->armors[0].header.frame_id;
        }
        else
        {
            std::cout << "unknown frame id!" << std::endl;
            return;
        }
        
        // 检查装甲板数量
        if(msg->num == 0)
        {
            std::cout << "Armors Empty" << std::endl; 
            m_armors_pub -> publish(armors_msg);
            return;
        }
        
        std::cout << "recieve " << msg->num << " armors" << std::endl;

        for(auto armor : msg -> armors)
        {
            base_interfaces::msg::Armor armor_msg;
            armor_msg.header     = armor.header;
            armor_msg.id         = armor.id;
            armor_msg.armor_type = armor.armor_type;
            armor_msg.color      = armor.color;
            armor_msg.vertices   = armor.vertices;

            // 验证数据正确性
            if (armor.armor_type < 0 || armor.armor_type > 2)
            {
                std::cout << "wrong armor type: " << armor.armor_type << std::endl;
                continue;
            }
            wmj::ARMORTYPE armor_type = wmj::ARMORTYPE(armor.armor_type);
            if (armor.vertices.size() != 4)
            {
                std::cout << "error: armor.vertices.size(): " << armor.vertices.size() << " is not 4" << std::endl;
                continue;
            }
            std::vector<cv::Point2f> armor_vertices;
            for(int j = 0 ; j < 4 ; j++)
            {
                cv::Point2f vertice;
                vertice.x = armor.vertices[j].position_x;
                vertice.y = armor.vertices[j].position_y;
                armor_vertices.push_back(vertice);
            }

            // 获取相机系下装甲板位姿和角度
            geometry_msgs::msg::Pose pose;
            double yaw_angle;
            if (getPositionAngle(armor_vertices, armor_type, armor.header.frame_id, pose, yaw_angle))
            {
                armor_msg.armor_pose = pose;
                armor_msg.yaw_angle = yaw_angle;
                
                armors_msg.armors.push_back(armor_msg);
                armors_msg.num++;
            }
            else
            {
                std::cout << "getPositionAngle error!" << std::endl;
                continue;
            }
        }

        m_armors_pub->publish(armors_msg);
        if (m_debug)
        {
            std::cout << "mono pub" << std::endl;
            double start_time, end_time;
            builtin_interfaces::msg::Time t = now();
            start_time = msg->armors[0].header.stamp.sec + msg->armors[0].header.stamp.nanosec * 1e-9;
            end_time = t.sec + t.nanosec * 1e-9;
            std::cout << "to mono time: " << end_time - start_time << std::endl;
        }
    }

    bool Monocular::getPositionAngle(std::vector<cv::Point2f> vertices,
                                wmj::ARMORTYPE armor_type,
                                std::string frame_id,
                                geometry_msgs::msg::Pose &pose,
                                double &yaw_angle)
    {
        std::vector<cv::Point2f> ipoints_v;
        ipoints_v.push_back(cv::Point2f(vertices[1].x, vertices[1].y));
        ipoints_v.push_back(cv::Point2f(vertices[0].x, vertices[0].y));
        ipoints_v.push_back(cv::Point2f(vertices[3].x, vertices[3].y));
        ipoints_v.push_back(cv::Point2f(vertices[2].x, vertices[2].y));

        // 寻找CameraMat下标
        std::vector<std::string>::iterator it = std::find(m_camera_frame_id.begin(), m_camera_frame_id.end(), frame_id);
        int index;
        if (it != m_camera_frame_id.end())
        {
            index = std::distance(m_camera_frame_id.begin(), it);
        }
        else{
            std::cout << "not find CameraMat of " << frame_id << "!" << std::endl;
            return false;
        }

        cv::Mat rvec, tvec, rmat;
        bool ret;   //solvePnP标志位

        // 判断装甲板类型
        if (armor_type == wmj::ARMORTYPE::ARMOR_LARGE)
        {
            ret = cv::solvePnP(m_obj_large, ipoints_v, m_CameraMat[index], m_DistMat[index], rvec, tvec, false, cv::SOLVEPNP_IPPE);
        }
        else if (armor_type == wmj::ARMORTYPE::ARMOR_SMALL)
        {
            ret = cv::solvePnP(m_obj_small, ipoints_v, m_CameraMat[index], m_DistMat[index], rvec, tvec, false, cv::SOLVEPNP_IPPE);
        }
        else{
            std::cout << "armor type is not small or large, armor type: " << armor_type << std::endl;
            return false;
        }
        if(!ret)
        {
            std::cout << "solvePnP wrong" <<std::endl;
            return false;
        }    
        cv::Rodrigues(rvec, rmat);  // 旋转向量转化旋转矩阵
        if (m_debug)
        {
            std::cout << "tvec: \n" << tvec << std::endl;
            // std::cout << "rmat: \n" << rmat << std::endl;
        }

        // 旋转矩阵转化为欧拉角
        Eigen::Matrix3d E_rmat;
        cv2eigen(rmat,E_rmat);
        Eigen::Quaterniond quaternion(E_rmat);

        // 设置位姿
        pose.position.x = tvec.ptr<double>(0)[0];
        pose.position.y = tvec.ptr<double>(0)[1];
        pose.position.z = tvec.ptr<double>(0)[2];
        pose.orientation.x = quaternion.x();
        pose.orientation.y = quaternion.y();
        pose.orientation.z = quaternion.z();
        pose.orientation.w = quaternion.w();

        if (m_debug)
        {
            std::cout << _red("distance: ") << _red(std::to_string(getDistance(cv::Point3d(pose.position.x, pose.position.y, pose.position.z)))) << std::endl;
        }

        // 计算偏航角
        double Z = asin(-1 * rmat.at<double>(1, 2));
        yaw_angle = asin(rmat.at<double>(1, 3) / cos(Z));   //yaw欧拉角，没用
        return true;
    }

    void Monocular::getObjPoints(float _width, float _height, wmj::ARMORTYPE armor_type)
    {
        float width = _width * 1., height = _height * 1.;

        // 判断装甲板类型
        if (armor_type == wmj::ARMORTYPE::ARMOR_LARGE)
        {
            m_obj_large.push_back(cv::Point3f(-width / 2.f, -height / 2.f, 0.0));
            m_obj_large.push_back(cv::Point3f(-width / 2.f, height / 2.f, 0.0));
            m_obj_large.push_back(cv::Point3f(width / 2.f, height / 2.f, 0.0));
            m_obj_large.push_back(cv::Point3f(width / 2.f, -height / 2.f, 0.0));
        }
        else if (armor_type == wmj::ARMORTYPE::ARMOR_SMALL)
        {
            m_obj_small.push_back(cv::Point3f(-width / 2.f, -height / 2.f, 0.0));
            m_obj_small.push_back(cv::Point3f(-width / 2.f, height / 2.f, 0.0));
            m_obj_small.push_back(cv::Point3f(width / 2.f, height / 2.f, 0.0));
            m_obj_small.push_back(cv::Point3f(width / 2.f, -height / 2.f, 0.0));
        }
        else
        {
            std::cout << "getObjPoints error, Unknown armor type: " << armor_type << std::endl;
        }
        return;
    }
   
}

int main(int argc, char* argv[])
{
    rclcpp::init(argc, argv);
    auto monocular = std::make_shared<wmj::Monocular>("monocular");
    rclcpp::spin(monocular);
    rclcpp::shutdown();
    return 0;
}